package com.dataflow.example.interceptor;

import cn.hutool.json.JSONUtil;
import com.etl.dataflow.executor.core.AbstractAfterWriteInterceptor;
import com.etl.dataflow.executor.core.TaskContext;
import com.etl.dataflow.executor.core.TaskContextHolder;
import com.etl.dataflow.executor.element.Column;
import com.etl.dataflow.executor.element.Row;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Scope;
import org.springframework.stereotype.Component;

import java.sql.Timestamp;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

@Slf4j
@Component("MyAfterWriteInterceptor")
@Scope("prototype")
public class MyAfterWriteInterceptor extends AbstractAfterWriteInterceptor {

    private final Map<String, Timestamp> maxIdtMap = new HashMap<>(8);

    @Override
    public void onAfterWrite(List<Row> writtenRows, List<Row> failedRows) {
//        log.debug(String.format("自定义AfterWrite拦截器已执行, 写入成功数量: %d，写入失败数量: %d.", writtenRows.size(), failedRows.size()));
        for (Row row : writtenRows) {
            Column<?> idtColumn = row.getColumnByName("IDT");
            Column<?> datasourceColumn = row.getColumnByName("DATA_SOURCES");
            if (datasourceColumn == null || idtColumn == null) {
                continue;
            }
            Timestamp timestamp = idtColumn.asTimestamp();
            String datasource = datasourceColumn.asString();
            if (!maxIdtMap.containsKey(datasource)) {
                maxIdtMap.put(datasource, new Timestamp(0));
            }
            if (timestamp.compareTo(maxIdtMap.get(datasource)) > 0) {
                maxIdtMap.replace(datasource, timestamp);
            }
        }
    }

    @Override
    public void onWriteFinished() {
        log.debug("自定义AfterWrite拦截器已执行, 全部数据已写入完毕，更新maxIdt...");
        saveMaxIdtToDb();
    }

    /**
     * 将各datasource的maxIdt持久化
     */
    public void saveMaxIdtToDb() {
        Optional<TaskContext> taskContext = TaskContextHolder.getTaskContext();
        if (taskContext.isPresent()) {
            TaskContext context = taskContext.get();
            String dstTableName = context.getTaskInfo().getTaskParam().getDstTableName();
            System.out.println(dstTableName + " ===>" + JSONUtil.toJsonStr(maxIdtMap));
        }
    }
}
